#!/usr/bin/env python3
# Author: Armit
# Create Time: 周二 2025/07/29 

# 最终解决方案 (提交用)

import json
from time import time
from pathlib import Path
from re import compile as Regex
from typing import Dict, Tuple

Loc = Tuple[Tuple[int, int]]
Annot = Dict[str, Loc]

BASE_PATH = Path(__file__).parent
DATA_PATH = BASE_PATH / 'data'
OUT_PATH = BASE_PATH / 'out'
OUT_PATH.mkdir(exist_ok=True)

PLT_NAME = 'plant'
CLS_NAME = 'class'
ENV_NAME = 'environment'
ARE_NAME = 'area'
ALT_NAME = 'altitude'

REGEX_CLS = Regex('([一二]级|[极濒易]危)')
REGEX_ALT = Regex('(海拔\d+米)')
REGEX_LOC = Regex('在(\w+?[省市区县村乡镇])[，。的中下]')
LOC_SUFFIX = list('省市区县村乡镇')


def predict_plant(text:str) -> Annot:
  pass

def predict_class(text:str) -> Annot:
  m = REGEX_CLS.findall(text)
  if not m: return
  key = m[0]
  loc = text.find(key)
  return {key: [[loc, loc + len(key) - 1]]}

def predict_environment(text:str) -> Annot:
  m = REGEX_LOC.findall(text)
  if not m: return
  key: str = m[0]
  if any(key.endswith(sfx) for sfx in LOC_SUFFIX): return
  loc = text.find(key)
  return {key: [[loc, loc + len(key) - 1]]}

def predict_area(text:str) -> Annot:
  m = REGEX_LOC.findall(text)
  if not m: return
  key: str = m[0]
  if not any(key.endswith(sfx) for sfx in LOC_SUFFIX): return
  loc = text.find(key)
  return {key: [[loc, loc + len(key) - 1]]}

def predict_altitude(text:str) -> Annot:
  m = REGEX_ALT.findall(text)
  if not m: return
  key = m[0]
  loc = text.find(key)
  return {key: [[loc, loc + len(key) - 1]]}


def run_predict():
  fp = DATA_PATH / 'test.jsonl'
  items = []
  with open(fp, 'r', encoding='utf-8') as fh:
    for line in fh.readlines():
      it = json.loads(line.strip())
      items.append(it)
  print(f'>> file {fp}')
  print(f'>> nlen: {len(items)}')

  PREDICTORS = {ptor_name: globals().get(f'predict_{ptor_name}') for ptor_name in [PLT_NAME, CLS_NAME, ENV_NAME, ARE_NAME, ALT_NAME]}
  PREDICT = []
  for it in items:
    text: str = it['text']
    label: Dict[str, Annot] = {}
    for ptor_name, predictor in PREDICTORS.items():
      pred = predictor(text)
      if pred is not None:
        label[ptor_name] = pred
    PREDICT.append({'text': text, 'label': label})

  save_fp = OUT_PATH / 'predict-submit.jsonl'
  print(f'>> save file: {save_fp}')
  with open(save_fp, 'w', encoding='utf-8') as fh:
    for pred in PREDICT:
      line = json.dumps(pred, indent=None, ensure_ascii=False)
      fh.write(line)
      fh.write('\n')


if __name__ == '__main__':
  ts_start = time()
  run_predict()
  ts_end = time()
  print(f'>> Done {ts_end - ts_start:.3f}s')
